import os
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.callbacks import CheckpointCallback
from env_turtlebot3 import TurtleBot3Env

def main():
    # 日志目录
    log_dir = "./ppo_tb3_log/"
    os.makedirs(log_dir, exist_ok=True)

    # 包装环境
    env = TurtleBot3Env()
    env = Monitor(env, log_dir)   # Monitor 负责记录 ep_len / ep_rew

    # 模型
    model = PPO(
        "MlpPolicy",
        env,
        verbose=1,
        tensorboard_log=log_dir   # 开启 TensorBoard 日志
    )

    # 保存模型的 callback，每隔1万步保存一次
    checkpoint_callback = CheckpointCallback(
        save_freq=10000,
        save_path="./ppo_tb3_checkpoints/",
        name_prefix="ppo_tb3"
    )

    # 开始训练
    model.learn(total_timesteps=50000, callback=checkpoint_callback)

    # 保存最终模型
    model.save("ppo_turtlebot3")

if __name__ == "__main__":
    main()

